# coding = utf-8

'''
基于densecrf的后处理
'''

import numpy as np
from scipy import ndimage
from skimage import measure
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import unary_from_labels, create_pairwise_bilateral, create_pairwise_gaussian
from PIL import Image

def select_tumor_densecrf(output, label, prob, images):
    predict_liver = np.zeros(output.shape)
    predict_tumor = np.zeros(output.shape)
    predict_liver[output>=1] = 1
    #predict_tumor[output == 2] = 1

    '''
       使用densecrf进行处理
    '''
    for i in range(images.shape[0]):
        image_item = images[i][1]
        prob_item = prob[i]
        image_item = image_item*255
        image_item = image_item.astype(np.uint8)
        image_gray = Image.fromarray(image_item).convert("L").convert("RGB")
        image_gray = np.array(image_gray)
        d = dcrf.DenseCRF(image_gray.shape[1] * image_gray.shape[0], 2)
        predict_1 = 1 - prob_item
        U = np.concatenate([predict_1.reshape(1, 512, 512), prob_item.reshape(1, 512, 512)], axis=0)
        U = -np.log(U)
        U = U.reshape((2, -1))
        d.setUnaryEnergy(U)
        feats = create_pairwise_gaussian(sdims=(3, 3), shape=(512, 512))
        d.addPairwiseEnergy(feats, compat=8, kernel=dcrf.DIAG_KERNEL,
                            normalization=dcrf.NORMALIZE_SYMMETRIC)

        feats = create_pairwise_bilateral(sdims=(10, 10), schan=(13, 13, 13),
                                          img=image_gray, chdim=2)
        d.addPairwiseEnergy(feats, compat=10,
                            kernel=dcrf.DIAG_KERNEL,
                            normalization=dcrf.NORMALIZE_SYMMETRIC)

        Q = d.inference(10)

        # 找出每个像素最可能的类
        MAP = np.argmax(Q, axis=0)
        MAP = MAP.reshape(image_item.shape)
        predict_tumor[i][MAP>0] = 1

    #select the biggest liver
    #predict_liver = ndimage.binary_dilation(predict_liver, iterations=1).astype(predict_liver.dtype)
    [liver_labels, num] = measure.label(predict_liver, return_num=True)
    region = measure.regionprops(liver_labels)
    box = []
    for i in range(num):
        box.append(region[i].area)
    label_num = box.index(max(box)) + 1
    liver_labels[liver_labels != label_num] = 0
    liver_labels[liver_labels == label_num] = 1
    #predict_liver = ndimage.binary_fill_holes(liver_labels).astype(int)

    #select tumor inside the liver
    label_copy = np.zeros(label.shape)
    label_copy[label >= 1] = 1
    predict_tumor = predict_tumor * label_copy


    #select tumor biggest
    real_tumor = np.zeros(predict_tumor.shape)

    #select tumor by prob each pixel
    '''
    prob = prob * label_copy
    real_tumor[prob>=0.6] = 1
    '''


    #predict_tumor = ndimage.binary_fill_holes(predict_tumor).astype(int)
    #select tumor by max prob

    [tumor_labels, num] = measure.label(predict_tumor, return_num=True)
    region = measure.regionprops(tumor_labels)
    box = []
    for i in range(num):
        tool_array = np.zeros(tumor_labels.shape)
        tool_array[tumor_labels == i+1] = 1
        tool_array = tool_array * prob
        box.append(np.max(tool_array))
        if np.max(np.array(tool_array)) >= 0.8:
            real_tumor[tumor_labels == i+1] =1
    print(sorted(box))


    '''
    label_copy = np.zeros(label.shape)
    label_copy[label == 2] = 1
    [_, num2] = measure.label(label_copy, return_num=True)
    print(num, num2, np.max(prob), np.min(prob))
    '''

    predict_liver[real_tumor == 1] = 2




    return predict_liver




if __name__ == '__main__':
    pass